// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/vector/bias_vector.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/constant_vector.h>
#include <pollux/vector/dictionary_vector.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/sequence_vector.h>
#include <pollux/vector/simple_vector.h>
#include <pollux/vector/vector_maker_stats.h>

namespace kumo::pollux {
    namespace detail {
        template<typename T>
        T jsonValue(const melon::Dynamic &jsonValue) {
            if constexpr (
                std::is_same_v<int8_t, T> || std::is_same_v<int16_t, T> ||
                std::is_same_v<int32_t, T> || std::is_same_v<int64_t, T>) {
                return jsonValue.asInt();
            }

            if constexpr (std::is_same_v<float, T> || std::is_same_v<double, T>) {
                return jsonValue.asDouble();
            }

            if constexpr (std::is_same_v<std::string, T>) {
                return jsonValue.asString();
            }

            if constexpr (std::is_same_v<bool, T>) {
                return jsonValue.asBool();
            }

            POLLUX_UNSUPPORTED();
        }

        template<typename T>
        void appendVariant(std::vector<variant> &values, const T &x) {
            values.push_back(x);
        };

        template<typename TupleT, std::size_t... Is>
        variant toVariantRow(const TupleT &tp, std::index_sequence<Is...>) {
            std::vector<variant> values;
            (appendVariant(values, std::get<Is>(tp)), ...);
            return variant::row(values);
        }

        template<typename TupleT, std::size_t TupleSize = std::tuple_size_v<TupleT> >
        variant toVariantRow(const TupleT &tp) {
            return toVariantRow(tp, std::make_index_sequence<TupleSize>{});
        }
    } // namespace detail

    class SimpleVectorLoader : public VectorLoader {
    public:
        explicit SimpleVectorLoader(std::function<VectorPtr(RowSet)> loader)
            : loader_(loader) {
        }

        void loadInternal(
            RowSet rows,
            ValueHook *hook,
            vector_size_t resultSize,
            VectorPtr *result) override {
            POLLUX_CHECK(!hook, "SimpleVectorLoader doesn't support ValueHook");
            auto &resultRef = *result;
            resultRef = loader_(rows);
            if (resultRef->size() < resultSize) {
                resultRef->resize(resultSize);
            }
        }

    private:
        std::function<VectorPtr(RowSet)> loader_;
    };

    class VectorMaker {
    public:
        explicit VectorMaker(memory::MemoryPool *pool) : pool_(pool) {
        }

        static std::function<bool(vector_size_t /*row*/)> null_every(
            int n,
            int startingFrom = 0) {
            return [n, startingFrom](vector_size_t row) {
                return row >= startingFrom && ((row - startingFrom) % n == 0);
            };
        }

        static std::shared_ptr<const RowType> row_type(
            std::vector<std::shared_ptr<const Type> > &&types);

        RowVectorPtr row_vector(const std::vector<VectorPtr> &children);

        RowVectorPtr row_vector(
            std::vector<std::string> childNames,
            const std::vector<VectorPtr> &children);

        RowVectorPtr row_vector(
            const std::shared_ptr<const RowType> &row_type,
            vector_size_t size);

        template<typename T>
        FlatVectorPtr<EvalType<T> > flat_vector(
            vector_size_t size,
            std::function<T(vector_size_t /*row*/)> value_at,
            std::function<bool(vector_size_t /*row*/)> is_null_at = nullptr,
            const TypePtr &type = CppToType<T>::create()) {
            auto flat_vector =
                    BaseVector::create<FlatVector<EvalType<T> > >(type, size, pool_);
            for (vector_size_t i = 0; i < size; i++) {
                if (is_null_at && is_null_at(i)) {
                    flat_vector->set_null(i, true);
                } else {
                    auto v = value_at(i);
                    flat_vector->set(i, EvalType<T>(v));
                }
            }
            return flat_vector;
        }

        template<typename T>
        std::shared_ptr<LazyVector> lazy_flat_vector(
            vector_size_t size,
            std::function<T(vector_size_t /*row*/)> value_at,
            std::function<bool(vector_size_t /*row*/)> is_null_at = nullptr) {
            return std::make_shared<LazyVector>(
                pool_,
                CppToType<T>::create(),
                size,
                std::make_unique<SimpleVectorLoader>([=](RowSet rowSet) {
                    // Populate requested rows with correct data and fill in gaps with
                    // "garbage".
                    SelectivityVector rows(rowSet.back() + 1, false);
                    for (auto row: rowSet) {
                        rows.setValid(row, true);
                    }
                    rows.updateBounds();

                    auto selectiveValueAt = [&](auto row) {
                        return rows.isValid(row) ? value_at(row) : T();
                    };

                    std::function<bool(vector_size_t)> selectiveIsNullAt = nullptr;
                    if (is_null_at) {
                        selectiveIsNullAt = [&](auto row) {
                            return rows.isValid(row) ? is_null_at(row) : false;
                        };
                    }

                    return flat_vector<T>(
                        rowSet.back() + 1, selectiveValueAt, selectiveIsNullAt);
                }));
        }

        template<typename T>
        FlatVectorPtr<T> flat_vector(
            size_t size,
            const TypePtr &type = CppToType<T>::create()) {
            return BaseVector::create<FlatVector<T> >(type, size, pool_);
        }

        /// Create a FlatVector<T>
        /// creates a FlatVector based on elements from the input std::vector. String
        /// vectors can be created using std::vector of char*, StringView,
        /// std::string, or std::string_view as input. The string contents will be
        /// copied to the flatvector's internal string buffer.
        ///
        /// Elements are non-nullable.
        ///
        /// Examples:
        ///   auto v1 = flat_vector({1, 2, 3, 4});
        ///   auto v2 = flat_vector({"hello", "world"});
        template<typename T>
        FlatVectorPtr<EvalType<T> > flat_vector(
            const std::vector<T> &data,
            const TypePtr &type = CppToType<T>::create());

        // Helper overload to allow users to use initializer list directly without
        // explicitly specifying the template type, e.g:
        //
        //   auto flatVector2 = flat_vector({"hello", "world"});
        template<typename T>
        FlatVectorPtr<EvalType<T> > flat_vector(
            const std::initializer_list<T> &data,
            const TypePtr &type = CppToType<T>::create()) {
            return flat_vector(std::vector<T>(data), type);
        }

        /// Create a FlatVector<T>
        /// creates a FlatVector based on elements from the input std::vector.
        /// Works for primitive and string types, similarly to flat_vector().
        ///
        /// Elements are nullable.
        ///
        /// Examples:
        ///   auto flat_vector = flat_vector_nullable({1, std::nullopt, 3});
        ///   auto flatVectorStr = flat_vector_nullable({
        ///       "hello"_sv, std::nullopt, "world"_sv});
        template<typename T>
        FlatVectorPtr<EvalType<T> > flat_vector_nullable(
            const std::vector<std::optional<T> > &data,
            const TypePtr &type = CppToType<T>::create());

        // Helper overload to allow users to use initializer list directly without
        // explicitly specifying the template type.
        template<typename T>
        FlatVectorPtr<EvalType<T> > flat_vector_nullable(
            const std::initializer_list<std::optional<T> > &data,
            const TypePtr &type = CppToType<T>::create()) {
            return flat_vector_nullable(std::vector<std::optional<T> >(data), type);
        }

        template<typename T, int TupleIndex, typename TupleType>
        FlatVectorPtr<T> flat_vector(
            const std::vector<TupleType> &data,
            const TypePtr &type) {
            auto vector = BaseVector::create<FlatVector<T> >(type, data.size(), pool_);
            for (vector_size_t i = 0; i < data.size(); ++i) {
                vector->set(i, std::get<TupleIndex>(data[i]));
            }
            return vector;
        }

        template<typename T>
        FlatVectorPtr<T> all_null_flat_vector(vector_size_t size) {
            auto flat_vector =
                    BaseVector::create<FlatVector<T> >(CppToType<T>::create(), size, pool_);
            for (vector_size_t i = 0; i < size; i++) {
                flat_vector->set_null(i, true);
            }
            return flat_vector;
        }

        /// Create a BiasVector<T>
        /// creates a BiasVector (vector encoded using bias encoding) based on a flat
        /// input from an std::vector.
        ///
        /// Elements are nullable.
        ///
        /// Example:
        ///   auto bias_vector = maker.bias_vector<int64_t>({10, 15, 13, 11, 12, 14});
        template<typename T>
        BiasVectorPtr<EvalType<T> > bias_vector(
            const std::vector<std::optional<T> > &data);

        /// Create a SequenceVector<T>
        /// creates a SequenceVector (vector encoded using RLE) based on a flat
        /// input from an std::vector.
        ///
        /// Elements are nullable.
        ///
        /// Example:
        ///   auto sequence_vector = maker.sequence_vector<int64_t>({
        ///       10, 10, 10, std::nullopt, 15, 15, std::nullopt, std::nullopt});
        template<typename T>
        SequenceVectorPtr<EvalType<T> > sequence_vector(
            const std::vector<std::optional<T> > &data);

        /// Create a ConstantVector<T>
        /// creates a ConstantVector (vector that represents a single constant value)
        /// based on a flat input from an std::vector. The input vector may contain
        /// several elements, but if the input vector contains more than one distinct
        /// element, it fails.
        ///
        /// Elements are nullable.
        ///
        /// Examples:
        ///   auto v1 = maker.constant_vector<int64_t>({11, 11, 11});
        ///   auto v2 = maker.constant_vector<int64_t>(
        ///        {std::nullopt, std::nullopt});
        template<typename T>
        ConstantVectorPtr<EvalType<T> > constant_vector(
            const std::vector<std::optional<T> > &data);

        /// Create a DictionaryVector<T>
        /// creates a dictionary encoded vector based on a flat input from an
        /// std::vector.
        ///
        /// Elements are nullable.
        ///
        /// Example:
        ///   auto dictionary_vector = maker.dictionary_vector<int64_t>({
        ///       10, 10, 10, std::nullopt, 15, 15, std::nullopt, std::nullopt});
        template<typename T>
        DictionaryVectorPtr<EvalType<T> > dictionary_vector(
            const std::vector<std::optional<T> > &data);

        /// Convenience function that creates an vector based on input std::vector
        /// data, encoded with given `vecType`.
        template<typename T>
        SimpleVectorPtr<EvalType<T> > encoded_vector(
            VectorEncoding::Simple vecType,
            const std::vector<std::optional<T> > &data) {
            switch (vecType) {
                case VectorEncoding::Simple::FLAT:
                    return flat_vector_nullable(data);
                case VectorEncoding::Simple::CONSTANT:
                    return constant_vector(data);
                case VectorEncoding::Simple::DICTIONARY:
                    return dictionary_vector(data);
                case VectorEncoding::Simple::SEQUENCE:
                    return sequence_vector(data);
                case VectorEncoding::Simple::BIASED:
                    return bias_vector(data);
                default:
                    POLLUX_UNSUPPORTED("Unsupported encoding type for VectorMaker.");
            }
            return nullptr;
        }

        /// Create a ArrayVector<T>
        /// size and null for individual array is determined by sizeAt and is_null_at
        /// value for individual array is determined by value_at.
        template<typename T>
        ArrayVectorPtr array_vector(
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
            std::function<T(vector_size_t /* idx */)> value_at,
            std::function<bool(vector_size_t /*row */)> is_null_at = nullptr,
            std::function<bool(vector_size_t /* idx */)> valueIsNullAt = nullptr,
            const TypePtr &arrayType = ARRAY(CppToType<T>::create())) {
            BufferPtr nulls;
            BufferPtr offsets;
            BufferPtr sizes;
            auto numElements =
                    create_offsets_and_sizes(size, sizeAt, is_null_at, &nulls, &offsets, &sizes);

            return std::make_shared<ArrayVector>(
                pool_,
                arrayType,
                nulls,
                size,
                offsets,
                sizes,
                flat_vector<T>(numElements, value_at, valueIsNullAt));
        }

        template<typename T>
        ArrayVectorPtr array_vector_impl(
            const TypePtr &type,
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
            std::function<T(vector_size_t /* row */, vector_size_t /* idx */)>
            value_at,
            std::function<bool(vector_size_t /*row */)> is_null_at = nullptr) {
            BufferPtr nulls;
            BufferPtr offsets;
            BufferPtr sizes;
            auto numElements =
                    create_offsets_and_sizes(size, sizeAt, is_null_at, &nulls, &offsets, &sizes);

            auto flat_vector = BaseVector::create<FlatVector<EvalType<T> > >(
                type->childAt(0), numElements, pool_);
            vector_size_t currentIndex = 0;
            for (vector_size_t i = 0; i < size; ++i) {
                if (is_null_at && is_null_at(i)) {
                    continue;
                }
                for (vector_size_t j = 0; j < sizeAt(i); ++j) {
                    auto value = value_at(i, j);
                    flat_vector->set(currentIndex, EvalType<T>(value));
                    currentIndex++;
                }
            }

            return std::make_shared<ArrayVector>(
                pool_,
                type,
                nulls,
                size,
                offsets,
                sizes,
                flat_vector,
                BaseVector::count_nulls(nulls, 0, size));
        }

        /// Create a ArrayVector<T>
        /// size and null for individual array is determined by sizeAt and is_null_at
        /// value for elements of each array in a given row is determined by value_at.
        template<typename T>
        ArrayVectorPtr array_vector(
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
            std::function<T(vector_size_t /* row */, vector_size_t /* idx */)>
            value_at,
            std::function<bool(vector_size_t /*row */)> is_null_at = nullptr,
            const TypePtr &arrayType = ARRAY(CppToType<T>::create())) {
            return array_vector_impl(arrayType, size, sizeAt, value_at, is_null_at);
        }

        template<typename T>
        ArrayVectorPtr array_vector_impl(
            const TypePtr &type,
            const std::vector<std::vector<T> > &data) {
            vector_size_t size = data.size();
            BufferPtr offsets = allocateOffsets(size, pool_);
            BufferPtr sizes = allocateSizes(size, pool_);

            auto rawOffsets = offsets->asMutable<vector_size_t>();
            auto rawSizes = sizes->asMutable<vector_size_t>();

            // Count number of elements.
            vector_size_t numElements = 0;
            for (const auto &array: data) {
                numElements += array.size();
            }

            // Create the underlying flat vector.
            auto flat_vector = BaseVector::create<FlatVector<EvalType<T> > >(
                type->childAt(0), numElements, pool_);

            vector_size_t currentIdx = 0;
            for (const auto &arrayValue: data) {
                *rawSizes++ = arrayValue.size();
                *rawOffsets++ = currentIdx;

                for (auto arrayElement: arrayValue) {
                    flat_vector->set(currentIdx++, EvalType<T>(arrayElement));
                }
            }

            return std::make_shared<ArrayVector>(
                pool_, type, nullptr, size, offsets, sizes, flat_vector);
        }

        /// Create a ArrayVector<T>
        /// array elements are created based on input std::vectors and are
        /// non-nullable.
        template<typename T>
        ArrayVectorPtr array_vector(
            const std::vector<std::vector<T> > &data,
            const TypePtr &elementType = CppToType<T>::create()) {
            return array_vector_impl(ARRAY(elementType), data);
        }

        /// Create an ArrayVector<ROW> from nested std::vectors of Variants.
        ArrayVectorPtr array_of_row_vector(
            const RowTypePtr &row_type,
            const std::vector<std::vector<variant> > &data);

        /// Creates an ARRAY(ROW(...)) vector from a list of lists of optional tuples.
        ///
        /// Allows to create arrays with null elements, but arrays themselves cannot
        /// be null. Members of the structs cannot be null either.
        ///
        /// Example:
        ///
        ///  std::vector<std::vector<std::optional<std::tuple<int32_t, std::string>>>>
        ///      data = {
        ///          {{{1, "red"}}, {{2, "blue"}}, {{3, "green"}}},
        ///          {},
        ///          {std::nullopt},
        ///          {{{4, "green"}}, {{5, "purple"}}},
        ///  };
        ///
        ///  auto array_vector = maker_.array_of_row_vector(data, ROW({INTEGER(),
        ///  VARCHAR()}));
        template<typename TupleT>
        ArrayVectorPtr array_of_row_vector(
            const std::vector<std::vector<std::optional<TupleT> > > &data,
            const RowTypePtr &row_type) {
            std::vector<std::vector<variant> > arrays;
            for (const auto &tuples: data) {
                std::vector<variant> elements;
                for (const auto &t: tuples) {
                    if (t.has_value()) {
                        elements.push_back(detail::toVariantRow(t.value()));
                    } else {
                        elements.push_back(variant::null(TypeKind::ROW));
                    }
                }
                arrays.push_back(elements);
            }

            return array_of_row_vector(row_type, arrays);
        }

        template<typename T>
        ArrayVectorPtr array_vector_nullable_impl(
            const TypePtr &type,
            const std::vector<std::optional<std::vector<std::optional<T> > > > &data) {
            POLLUX_CHECK(type->is_array(), "Type must be an array: {}", type->toString());

            vector_size_t size = data.size();
            BufferPtr offsets = AlignedBuffer::allocate<vector_size_t>(size, pool_);
            BufferPtr sizes = AlignedBuffer::allocate<vector_size_t>(size, pool_);
            BufferPtr nulls = AlignedBuffer::allocate<uint64_t>(size, pool_);

            auto rawOffsets = offsets->asMutable<vector_size_t>();
            auto rawSizes = sizes->asMutable<vector_size_t>();
            auto rawNulls = nulls->asMutable<uint64_t>();
            bits::fillBits(rawNulls, 0, size, pool_);

            // Count number of elements.
            vector_size_t numElements = 0;
            vector_size_t indexPtr = 0;
            for (const auto &array: data) {
                numElements += array.has_value() ? array.value().size() : 0;
                if (!array.has_value()) {
                    bits::set_null(rawNulls, indexPtr, true);
                }
                indexPtr++;
            }

            using V = typename CppToType<T>::NativeType;

            // Create the underlying flat vector.
            auto flat_vector =
                    BaseVector::create<FlatVector<V> >(type->childAt(0), numElements, pool_);
            auto elementRawNulls = flat_vector->mutable_raw_nulls();

            vector_size_t currentIdx = 0;

            for (const auto &arrayValue: data) {
                *rawSizes++ = arrayValue.has_value() ? arrayValue.value().size() : 0;
                *rawOffsets++ = currentIdx;

                if (arrayValue.has_value()) {
                    for (auto arrayElement: arrayValue.value()) {
                        if (arrayElement == std::nullopt) {
                            bits::set_null(elementRawNulls, currentIdx, true);
                        } else {
                            flat_vector->set(currentIdx, V(*arrayElement));
                        }
                        ++currentIdx;
                    }
                }
            }

            return std::make_shared<ArrayVector>(
                pool_, type, nulls, size, offsets, sizes, flat_vector);
        }

        /// Create a ArrayVector<T>
        /// array elements are created based on input std::vectors and are
        /// nullable.
        template<typename T>
        ArrayVectorPtr array_vector_nullable(
            const std::vector<std::optional<std::vector<std::optional<T> > > > &data,
            const TypePtr &type = ARRAY(CppToType<T>::create())) {
            return array_vector_nullable_impl(type, data);
        }

        /// Creates an ArrayVector from a list of JSON arrays.
        ///
        /// JSON arrays can represent a null array, an empty array or array with null
        /// elements.
        ///
        /// Examples:
        ///  [1, 2, 3]
        ///  [1, 2, null, 4]
        ///  [null, null]
        ///  [] - empty array
        ///  null - null array
        ///
        /// @tparam T Type of array elements.
        /// @param jsonArrays A list of JSON arrays. JSON array cannot be an empty
        /// string.
        template<typename T>
        ArrayVectorPtr array_vector_from_json(
            const std::vector<std::string> &jsonArrays,
            const TypePtr &arrayType = ARRAY(CppToType<T>::create())) {
            std::vector<std::optional<std::vector<std::optional<T> > > > arrays;
            for (const auto &jsonArray: jsonArrays) {
                POLLUX_CHECK(!jsonArray.empty());

                const melon::Dynamic arrayObject = melon::parseJson(jsonArray);
                if (arrayObject.isNull()) {
                    // Null array.
                    arrays.push_back(std::nullopt);
                    continue;
                }
                std::vector<std::optional<T> > elements;
                append_elements_from_json_array(arrayObject, elements);
                arrays.push_back(elements);
            }
            return array_vector_nullable<T>(arrays, arrayType);
        }

        /// Creates an ArrayVector from a list of JSON arrays of arrays.
        ///
        /// JSON array of arrays can represent a null array, an empty array or array
        /// with null elements.
        ///
        /// Examples:
        /// [[1, 2], [2, 3, 4], [null, 7]]
        /// [[1, 3, 7, 9], []]
        /// [] - empty array of arrays
        /// null - null array of arrays
        /// [null] - array of null array
        /// [[]]
        ///
        /// @tparam T Type of array elements.
        /// @param jsonArrays A list of JSON arrays. JSON array cannot be an empty
        /// string.
        template<typename T>
        ArrayVectorPtr nested_array_vector_from_json(
            const std::vector<std::string> &jsonArrays,
            const TypePtr &arrayType = ARRAY(CppToType<T>::create())) {
            std::vector<std::optional<std::vector<std::optional<T> > > > baseVector;
            std::vector<vector_size_t> offsets;
            std::vector<vector_size_t> nulls;
            const std::vector<std::optional<T> > empty;
            int offset = 0;
            for (auto i = 0; i < jsonArrays.size(); i++) {
                const auto &jsonArray = jsonArrays.at(i);
                POLLUX_CHECK(!jsonArray.empty());
                const melon::Dynamic arraysObject = melon::parseJson(jsonArray);
                offsets.push_back(offset);
                if (arraysObject.isNull()) {
                    // Null array.
                    nulls.push_back(i);
                    continue;
                }
                for (const auto &nestedArray: arraysObject) {
                    if (nestedArray.isNull()) {
                        // Null nested array
                        baseVector.push_back(std::nullopt);
                    } else {
                        std::vector<std::optional<T> > elements;
                        append_elements_from_json_array(nestedArray, elements);
                        baseVector.push_back(elements);
                    }
                }
                offset += arraysObject.size();
            }
            auto baseArrayVector = array_vector_nullable<T>(baseVector, arrayType);
            return array_vector(offsets, baseArrayVector, nulls);
        }

        ArrayVectorPtr all_null_array_vector(
            vector_size_t size,
            const TypePtr &elementType);

        /// Create a Map<TKey, TValue>
        /// size and null for individual map is determined by sizeAt and is_null_at
        /// key and value for individual map is determined by keyAt and value_at
        template<typename TKey, typename TValue>
        MapVectorPtr map_vector(
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
            std::function<TKey(vector_size_t /* idx */)> keyAt,
            std::function<TValue(vector_size_t /* idx */)> value_at,
            std::function<bool(vector_size_t /*row */)> is_null_at = nullptr,
            std::function<bool(vector_size_t /*row */)> valueIsNullAt = nullptr,
            const TypePtr &type =
                    MAP(CppToType<TKey>::create(), CppToType<TValue>::create())) {
            BufferPtr nulls;
            BufferPtr offsets;
            BufferPtr sizes;
            auto numElements =
                    create_offsets_and_sizes(size, sizeAt, is_null_at, &nulls, &offsets, &sizes);

            return std::make_shared<MapVector>(
                pool_,
                type,
                nulls,
                size,
                offsets,
                sizes,
                flat_vector<TKey>(numElements, keyAt, nullptr, type->childAt(0)),
                flat_vector<TValue>(
                    numElements, value_at, valueIsNullAt, type->childAt(1)),
                BaseVector::count_nulls(nulls, 0, size));
        }

        template<typename TKey, typename TValue>
        MapVectorPtr map_vector(
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* mapRow */)> sizeAt,
            std::function<TKey(vector_size_t /* mapRow */, vector_size_t /*row*/)>
            keyAt,
            std::function<TValue(vector_size_t /* mapRow */, vector_size_t /*row*/)>
            value_at,
            std::function<bool(vector_size_t /*mapRow */)> is_null_at = nullptr) {
            BufferPtr nulls;
            BufferPtr offsets;
            BufferPtr sizes;
            auto numElements =
                    create_offsets_and_sizes(size, sizeAt, is_null_at, &nulls, &offsets, &sizes);

            auto rawNulls = nulls ? nulls->asMutable<uint64_t>() : nullptr;
            auto rawSizes = sizes->asMutable<vector_size_t>();

            std::vector<TKey> keys;
            keys.reserve(numElements);
            std::vector<TValue> values;
            values.reserve(numElements);
            for (vector_size_t mapRow = 0; mapRow < size; mapRow++) {
                if (rawNulls && bits::isBitNull(rawNulls, mapRow)) {
                    continue;
                }

                auto mapSize = rawSizes[mapRow];
                for (vector_size_t row = 0; row < mapSize; row++) {
                    keys.push_back(keyAt(mapRow, row));
                    values.push_back(value_at(mapRow, row));
                }
            }

            return std::make_shared<MapVector>(
                pool_,
                MAP(CppToType<TKey>::create(), CppToType<TValue>::create()),
                nulls,
                size,
                offsets,
                sizes,
                flat_vector(keys),
                flat_vector(values));
        }

        template<typename TKey, typename TValue>
        MapVectorPtr map_vector(
            const std::vector<std::vector<std::pair<TKey, std::optional<TValue> > > > &
            maps,
            const TypePtr &mapType =
                    MAP(CppToType<TKey>::create(), CppToType<TValue>::create())) {
            std::vector<
                        std::optional<std::vector<std::pair<TKey, std::optional<TValue> > > > >
                    nullableMaps;
            nullableMaps.reserve(maps.size());
            for (const auto &m: maps) {
                nullableMaps.push_back(m);
            }

            return map_vector<TKey, TValue>(nullableMaps, mapType, false);
        }

        template<typename TKey, typename TValue>
        MapVectorPtr map_vector(
            const std::vector<std::optional<
                std::vector<std::pair<TKey, std::optional<TValue> > > > > &maps,
            const TypePtr &mapType =
                    MAP(CppToType<TKey>::create(), CppToType<TValue>::create()),
            bool hasNulls = true) {
            std::vector<TKey> keys;
            std::vector<TValue> values;
            std::vector<bool> nullValues;

            for (const auto &map: maps) {
                if (map.has_value()) {
                    for (const auto &[key, value]: map.value()) {
                        keys.push_back(key);
                        values.push_back(value.value_or(TValue()));
                        nullValues.push_back(!value.has_value());
                    }
                }
            }

            std::function<bool(vector_size_t)> is_null_at = nullptr;
            if (hasNulls) {
                is_null_at = [&](vector_size_t row) { return !maps[row].has_value(); };
            }

            return map_vector<TKey, TValue>(
                maps.size(),
                [&](vector_size_t row) {
                    return maps[row].has_value() ? maps[row]->size() : 0;
                },
                [&](vector_size_t idx) { return keys[idx]; },
                [&](vector_size_t idx) { return values[idx]; },
                is_null_at,
                [&](vector_size_t idx) { return nullValues[idx]; },
                mapType);
        }

        /// Creates a MapVector from a list of JSON maps.
        ///
        /// JSON maps can represent a null map, an empty map or a map with null
        /// values. Null keys are not allowed.
        ///
        /// Note that order of map keys in the MapVector is not guaranteed to match
        /// the order of map keys in the JSON.
        ///
        /// Examples:
        ///  {1: 10, 2: 20, 3: 30}
        ///  {1: 10, 2: 20, 3: null, 4: 40}
        ///  {1: null, 2: null}
        ///  {} - empty map
        ///  null - null map
        ///
        /// @tparam K Type of map keys. Must be a std::string or an integer: int8_t,
        /// int16_t, int32_t, int64_t.
        /// @tparam V Type of map value. Can be an integer or a floating point
        /// number.
        /// @param jsonMaps A list of JSON maps. JSON map cannot be an empty
        /// string.
        template<typename K, typename V>
        MapVectorPtr map_vector_from_json(
            const std::vector<std::string> &jsonMaps,
            const TypePtr &mapType =
                    MAP(CppToType<K>::create(), CppToType<V>::create())) {
            static_assert(
                std::is_same_v<K, int8_t> || std::is_same_v<K, int16_t> ||
                std::is_same_v<K, int32_t> || std::is_same_v<K, int64_t> ||
                std::is_same_v<K, float> || std::is_same_v<K, double> ||
                std::is_same_v<K, std::string>);

            std::vector<std::optional<std::vector<std::pair<K, std::optional<V> > > > >
                    maps;
            for (const auto &jsonMap: jsonMaps) {
                POLLUX_CHECK(!jsonMap.empty());

                melon::json::serialization_opts options;
                options.convert_int_keys = true;
                options.allow_non_string_keys = true;
                options.allow_nan_inf = true;
                melon::Dynamic mapObject = melon::parseJson(jsonMap, options);
                if (mapObject.isNull()) {
                    // Null map.
                    maps.push_back(std::nullopt);
                    continue;
                }

                std::vector<std::pair<K, std::optional<V> > > pairs;
                for (const auto &item: mapObject.items()) {
                    auto key = detail::jsonValue<K>(item.first);

                    if (item.second.isNull()) {
                        // Null value.
                        pairs.push_back({key, std::nullopt});
                    } else {
                        pairs.push_back({key, detail::jsonValue<V>(item.second)});
                    }
                }

                maps.push_back(pairs);
            }

            return map_vector<K, V>(maps, mapType);
        }

        MapVectorPtr all_null_map_vector(
            vector_size_t size,
            const TypePtr &keyType,
            const TypePtr &valueType);

        /// Create a FlatVector from a variant containing a scalar value.
        template<TypeKind kind>
        VectorPtr to_flat_vector(variant value) {
            using T = typename TypeTraits<kind>::NativeType;
            if constexpr (std::is_same_v<T, StringView>) {
                return flat_vector<StringView>({StringView(value.value<const char *>())});
            } else {
                return flat_vector(std::vector<T>(1, value.value<T>()));
            }
        }

        /// Create constant vector of type ROW from a variant.
        VectorPtr
        constant_row(const RowTypePtr &row_type, variant value, vector_size_t size) {
            POLLUX_CHECK_EQ(value.kind(), TypeKind::ROW);

            std::vector<VectorPtr> fields(row_type->size());
            for (auto i = 0; i < row_type->size(); i++) {
                fields[i] = POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                            to_flat_vector, row_type->childAt(i)->kind(), value.row()[i]);
            }

            return BaseVector::wrap_in_constant(
                size,
                0,
                std::make_shared<RowVector>(
                    pool_, row_type, nullptr, 1, std::move(fields)));
        }

        template<typename T = BaseVector>
        static std::shared_ptr<T> flatten(const VectorPtr &vector) {
            SelectivityVector allRows(vector->size());
            auto flat_vector =
                    BaseVector::create<T>(vector->type(), vector->size(), vector->pool());
            flat_vector->copy(vector.get(), allRows, nullptr);
            return flat_vector;
        }

        /// Create an ArrayVector from a vector of offsets and a base element
        /// vector. The size of the arrays is computed from the difference of
        /// offsets. An optional vector of nulls can be passed to specify null rows.
        /// The offset for a null value must match previous offset
        /// i.e size computed should be zero.
        /// E.g array_vector({0, 2 ,2}, elements, {1}) creates an array vector
        /// with array at index 1 as null.
        ArrayVectorPtr array_vector(
            const std::vector<vector_size_t> &offsets,
            const VectorPtr &elements,
            const std::vector<vector_size_t> &nulls = {});

        /// Create a MapVector from a vector of offsets and key and value vectors.
        /// The size of the maps is computed from the difference of offsets.
        /// The sizes of the key and value vectors must be equal.
        /// An optional vector of nulls can be passed to specify null rows.
        /// The offset for a null value must match previous offset
        /// i.e size computed should be zero.
        /// E.g map({0, 2 ,2}, keys, values, {1}) creates a map vector
        /// with map at index 1 as null.
        MapVectorPtr map_vector(
            const std::vector<vector_size_t> &offsets,
            const VectorPtr &keys,
            const VectorPtr &values,
            const std::vector<vector_size_t> &nulls = {});

    private:
        vector_size_t create_offsets_and_sizes(
            vector_size_t size,
            std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
            std::function<bool(vector_size_t /*row */)> is_null_at,
            BufferPtr *nulls,
            BufferPtr *offsets,
            BufferPtr *sizes);

        template<typename T>
        void append_elements_from_json_array(
            const melon::Dynamic &arrayObject,
            std::vector<std::optional<T> > &elements) {
            for (const auto &element: arrayObject) {
                if (element.isNull()) {
                    elements.push_back(std::nullopt);
                } else {
                    elements.push_back(detail::jsonValue<T>(element));
                }
            }
        }

        memory::MemoryPool *pool_;
    };
} // namespace kumo::pollux

#include <pollux/vector/vector_maker-inl.h>
